import os
os.environ["OPENAI_API_KEY"] = 'sk-ce243630991443b2a32a217c15f8128c'
os.environ["OPENAI_BASE_URL"] = 'https://dashscope.aliyuncs.com/compatible-mode/v1'

# 1.Load 导入Document Loaders
from langchain_community.document_loaders import PyPDFLoader
from langchain_community.document_loaders import Docx2txtLoader
from langchain_community.document_loaders import TextLoader

# 加载Documents
base_dir = '/Users/shenhao/Documents/PycharmProjects/langchain-in-action/02_文档QA系统/OneFlower' # 文档的存放目录
documents = []
for file in os.listdir(base_dir):
    # 构建完整的文件路径
    file_path = os.path.join(base_dir, file)
    if file.endswith('.pdf'):
        loader = PyPDFLoader(file_path)
        documents.extend(loader.load())
    elif file.endswith('.docx'):
        loader = Docx2txtLoader(file_path)
        documents.extend(loader.load())
    elif file.endswith('.txt'):
        loader = TextLoader(file_path)
        documents.extend(loader.load())

# 2.Split 将Documents切分成块以便后续进行嵌入和向量存储
from langchain.text_splitter import RecursiveCharacterTextSplitter
text_splitter = RecursiveCharacterTextSplitter(chunk_size=200, chunk_overlap=10)
chunked_documents = text_splitter.split_documents(documents)

# 3.Store 将分割嵌入并存储在矢量数据库Qdrant中
from langchain_community.vectorstores import Qdrant
from langchain_openai import OpenAIEmbeddings
from langchain_dashscope import DashScopeEmbeddings

# 初始化嵌入模型
embeddings = DashScopeEmbeddings(
    model="text-embedding-v4",  # 可选：text-embedding-v1, text-embedding-v2
    api_key="sk-ce243630991443b2a32a217c15f8128c",
)
# vectorstore = Qdrant.from_documents(
#     documents=chunked_documents, # 以分块的文档
#     embedding=embeddings, # 用OpenAI的Embedding Model做嵌入
#     location=":memory:",  # in-memory 存储
#     collection_name="my_documents",) # 指定collection_name
# text-embedding-v4模型最大支持10条，分批处理文档，每批不超过10条
batch_size = 10
vectorstore = None

for i in range(0, len(chunked_documents), batch_size):
    batch_documents = chunked_documents[i:i + batch_size]

    if vectorstore is None:
        # 第一批文档创建向量存储
        vectorstore = Qdrant.from_documents(
            documents=batch_documents,
            embedding=embeddings,
            location=":memory:",
            collection_name="my_documents",
        )
    else:
        # 后续批次添加到现有向量存储
        vectorstore.add_documents(documents=batch_documents)

# 4. Retrieval 准备模型和Retrieval链
import logging # 导入Logging工具
from langchain_openai import ChatOpenAI # ChatOpenAI模型
from langchain.retrievers.multi_query import MultiQueryRetriever # MultiQueryRetriever工具
from langchain.chains import RetrievalQA # RetrievalQA链

# 设置Logging
logging.basicConfig()
logging.getLogger('langchain.retrievers.multi_query').setLevel(logging.INFO)

# 实例化一个大模型工具 - OpenAI的GPT-3.5
llm = ChatOpenAI(model="deepseek-v3", temperature=0)

# 实例化一个MultiQueryRetriever
# 1.MultiQueryRetriever在代码中的意义是什么？
# MultiQueryRetriever - 这是LangChain中的检索器之一。它的特点使用语言模型从不同的视角为给定的用户输入查询生成多个查询。然后使用这些查询来检索一组相关文档。—— 这是一种常见策略。
# 2.Langchain中不同Retrievers应该如何选择？——  Langchain中Retriever很多，MultiQueryRetriever、VectorRetriever和KeywordRetriever等等。详见第15课。
# 不同的Retrievers在Langchain中有不同的选择方式。选择Retriever的方法取决于您的需求和数据类型。
# 如果数据是基于向量的，可以选择使用VectorRetriever。VectorRetriever使用向量数据库进行检索，并根据查询向量找到最相似的文档。可以使用不同的嵌入模型（如OpenAIEmbeddings）来构建向量数据库，并使用VectorRetriever进行检索。
# 如果数据是基于关键词的，可以选择使用KeywordRetriever。KeywordRetriever使用关键词匹配的方法进行检索。可以将关键词与文档进行关联，并使用KeywordRetriever根据关键词进行检索。
# 如果希望自动生成多个查询并获取更丰富的结果，可以选择使用MultiQueryRetriever。MultiQueryRetriever使用语言模型生成多个查询，并根据这些查询检索相关文档。MultiQueryRetriever生成多个查询的不同视角，以获取更多潜在相关的文档。
retriever_from_llm = MultiQueryRetriever.from_llm(retriever=vectorstore.as_retriever(), llm=llm)

# 实例化一个RetrievalQA链
qa_chain = RetrievalQA.from_chain_type(llm,retriever=retriever_from_llm)

# 5. Output 问答系统的UI实现
from flask import Flask, request, render_template

# app = Flask(__name__)  # Flask APP
# app = Flask(__name__, template_folder="/templates/index.html")  # 明确指定模板目录
project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
template_dir = os.path.join(project_root, 'templates')

app = Flask(__name__, template_folder=template_dir)


@app.route('/', methods=['GET', 'POST'])
def home():
    if request.method == 'POST':
        # 接收用户输入作为问题
        question = request.form.get('question')

        # RetrievalQA链 - 读入问题，生成答案
        result = qa_chain({"query": question})

        # 把大模型的回答结果返回网页进行渲染
        return render_template('index.html', result=result)

    return render_template('index.html')


if __name__ == "__main__":
    app.run(host='0.0.0.0', debug=True, port=5001)